import os

import numpy as np
import geopandas as gpd
import matplotlib.pyplot as plt
from PIL import Image
from fontTools.config import Config
from matplotlib.colors import ListedColormap, BoundaryNorm, LinearSegmentedColormap
from pykrige import UniversalKriging
import rasterio
from rasterio.mask import mask
import json
from scipy.ndimage import gaussian_filter
from shapely.geometry.geo import box
import yaml
import time
import datetime
def read_json(file_path):
    """读取 JSON 文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def load_config(path = "application.yaml"):
    with open(path, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
        return data

""" 获取保存文件名称 """
def get_file_name(ext):
    base_path = CONFIG['result_img']['base_path']
    if not os.path.exists(base_path):
        # 文件夹不存在，创建文件夹
        os.makedirs(base_path)
        print(f"文件夹 '{base_path}' 已创建。")
    file_name_template = CONFIG['result_img']['file_name_template']
    current_date = datetime.datetime.now()
    file_name_date = current_date.strftime(file_name_template)
    return base_path + file_name_date + ext

def create_grid(bounding_box, grid_size=1024):
    """生成插值网格"""
    grid_lon = np.linspace(bounding_box["bottom_left"][0], bounding_box["top_right"][0], grid_size)
    grid_lat = np.linspace(bounding_box["bottom_left"][1], bounding_box["top_right"][1], grid_size)
    return np.meshgrid(grid_lon, grid_lat)


def perform_kriging(longitude, latitude, values, grid_lon, grid_lat):
    """使用PyKrige进行插值"""
    krige = UniversalKriging(longitude, latitude, values, variogram_model='linear',  drift_terms=["regional_linear"],)
    # krige = OrdinaryKriging(longitude, latitude, values, variogram_model='linear')
    return krige.execute("grid", grid_lon, grid_lat)[0]


def load_geojson(geojson_path):
    """加载并确保 GeoJSON 文件的几何在插值网格范围内"""
    geojson = gpd.read_file(geojson_path)
    geojson = geojson.to_crs(epsg=4326)  # 投影为 WGS84
    return geojson.geometry.union_all()


def check_overlap(region, bounding_box):
    """检查 GeoJSON 区域与插值网格的范围是否有重叠"""
    bbox_polygon = gpd.GeoSeries([box(
        bounding_box["bottom_left"][0], bounding_box["bottom_left"][1],
        bounding_box["top_right"][0], bounding_box["top_right"][1]
    )], crs="EPSG:4326")

    if not region.intersects(bbox_polygon.union_all()):
        raise ValueError("GeoJSON 区域与插值网格范围无重叠，请检查输入数据")


def smooth_data(diff_aqi, sigma=2):
    """高斯平滑"""
    return gaussian_filter(diff_aqi, sigma)

def crop_with_geojson(smoothed_diff_aqi, smoothed_region, transform):
    """使用 GeoJSON 边界裁剪数据，并将区域外的像素值设置为白色"""
    with rasterio.io.MemoryFile() as memfile:
        with memfile.open(
                driver='GTiff',
                height=smoothed_diff_aqi.shape[0],
                width=smoothed_diff_aqi.shape[1],
                count=1,
                crs='EPSG:4326',
                dtype='float32',
                transform=transform
        ) as dataset:
            dataset.write(smoothed_diff_aqi, 1)
            # 执行裁剪
            out_image, out_transform = mask(dataset, [smoothed_region], crop=True, all_touched=True)

            # 将区域外的颜色设置为白色
            # 假设白色用NaN表示，可以根据需求修改为0或其他值
            out_image[out_image == 0] = np.nan  # 如果区域外的值是0，将其替换为NaN（或白色）
            return out_image, out_transform

def plot_data(resized_array, bounding_box,norm=None,cmap = "coolwarm"):
    """绘制结果并保存为SVG格式"""
    plt.figure(figsize=(10, 6))
    plt.axis('off')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    plt.imshow(resized_array, extent=(
        bounding_box["bottom_left"][0], bounding_box["bottom_left"][1],
        bounding_box["top_right"][0], bounding_box["top_right"][1]
    ),
               origin="upper", cmap=cmap,norm=norm,interpolation="bilinear")
    plt.savefig(get_file_name(CONFIG['result_img']['ext_name']), dpi=1000, pad_inches=0, bbox_inches='tight',transparent=True)
    plt.show()


CONFIG = load_config("application.yaml")
def main():
    """主程序执行"""
    bounding_box = {
        "bottom_left": (CONFIG['bounding_box'][0],CONFIG['bounding_box'][1]),  # 左下角坐标 , 3.513421, 73.388672
        "top_right": (CONFIG['bounding_box'][2],CONFIG['bounding_box'][3]),  # 右上角坐标  53.696706, 136.142578
    }
    geojson_path = CONFIG['geo_json']
    json_file = "data.json"
    colors = CONFIG['colors']
    bounds = CONFIG['bounds']  # 定义值的区间

    # 1. 加载 JSON 数据
    data = read_json(json_file)
    longitude = [point[0] for point in data]
    latitude = [point[1] for point in data]
    values = [point[2] for point in data]

    # 2. 创建网格
    grid_x, grid_y = create_grid(bounding_box,CONFIG['dpi'])
    # 3. 插值计算
    grid_aqi = perform_kriging(longitude, latitude, values, grid_x[0], grid_y[:, 0])

    cmap = ListedColormap(colors)
    # cmap = LinearSegmentedColormap.from_list("my_gradient", colors)
    norm = BoundaryNorm(bounds, cmap.N)
    print(f"数据值范围: {min(values)} ~ {max(values)}")
    print(f"插值结果范围: {grid_aqi.min()} ~ {grid_aqi.max()}")
    print(f"颜色区间: {bounds}")

    # 4. 计算差值
    diff_aqi = grid_aqi - np.mean(values)

    # 5. 加载并检查GeoJSON
    geojson_region = load_geojson(geojson_path)
    check_overlap(geojson_region, bounding_box)

    # 6. 创建裁剪掩膜
    transform = rasterio.transform.from_bounds(
        bounding_box["bottom_left"][0], bounding_box["bottom_left"][1],
        bounding_box["top_right"][0], bounding_box["top_right"][1],
        grid_aqi.shape[1], grid_aqi.shape[0]
    )

    # 7. 高斯平滑
    # smoothed_diff_aqi = smooth_data(diff_aqi)

    # 8. 对GeoJSON边界平滑
    # smoothed_region = geojson_region.buffer(0.001)

    # 9. 裁剪数据
    out_image, out_transform = crop_with_geojson(grid_aqi, geojson_region, transform)

    # 10. 检查裁剪结果
    # if out_image is None or np.all(out_image == 0):
    #     raise ValueError("裁剪后结果为空，请检查 GeoJSON 数据与插值网格的范围匹配")

    # 11. 图像处理
    # image = Image.fromarray(out_image[0])
    # resized_image = image.resize((image.width * 2, image.height * 2), Image.LANCZOS)
    # resized_array = np.array(resized_image)
    # 12. 绘图
    plot_data(out_image[0], bounding_box,norm,cmap)


if __name__ == "__main__":
    start = time.time()
    main()
    end = time.time()
    print(f"Process Run {end - start} Seconds")
